# -*- coding: utf-8 -*-
"""
@Time    : 2024/9/5 11:23 
@Author  : ZhangShenao 
@File    : main.py 
@Desc    : 多Agent执行入口
"""
from langchain_core.runnables import RunnableConfig
from langgraph.graph import StateGraph

from base import BaseAgentState
from live_agent import create_live_graph
from xhs_agent import create_xhs_graph


# 定义入口节点
def entry_node(state: BaseAgentState, config: RunnableConfig) -> BaseAgentState:
    """入口节点"""

    # 透传并返回状态
    return state


if __name__ == '__main__':
    # 创建整体图架构,并指定状态类型为BaseAgentState
    graph_builder = StateGraph(BaseAgentState)

    # 创建live_agent和xhs_agent两个子图
    live_graph = create_live_graph()
    xhs_graph = create_xhs_graph()

    # 添加节点
    graph_builder.add_node(node="entry", action=entry_node)
    graph_builder.add_node(node="live_agent", action=live_graph)
    graph_builder.add_node(node="xhs_agent", action=xhs_graph)

    # 添加边,两个子图为两个并行边
    graph_builder.set_entry_point(key="entry")
    graph_builder.add_edge(start_key="entry", end_key="live_agent")
    graph_builder.add_edge(start_key="entry", end_key="xhs_agent")
    # 两个子图均为结束节点
    graph_builder.set_finish_point(key="live_agent")
    graph_builder.set_finish_point(key="xhs_agent")

    # 编译整个Graph结构,生成Runnable可运行组件
    graph = graph_builder.compile()

    # 执行Graph,获取结果
    result = graph.invoke({"query": "潮汕牛肉丸"})
    print(result)
